import numpy as np

from htm.Connection import Connection


class ColumnPoolerRegion:

    def __init__(self, columnPoolerRegionInit, lateralCorColOrders):
        # Parameters
        self.inputWidth = columnPoolerRegionInit['inputWidth']
        self.cellCount = columnPoolerRegionInit['cellCount']
        self.sdrSize = columnPoolerRegionInit['sdrSize']
        self.onlineLearning = columnPoolerRegionInit['onlineLearning']
        self.maxSdrSize = columnPoolerRegionInit['maxSdrSize']
        self.minSdrSize = columnPoolerRegionInit['minSdrSize']
        self.predictedInhibitionThreshold = columnPoolerRegionInit['predictedInhibitionThreshold']
        self.inertiaFactor = columnPoolerRegionInit['inertiaFactor']
        # Initialization
        self.learn = True
        self.useInertia = True
        self.objectL2RepresentationsMatrix = np.zeros([0, self.cellCount])
        self.proximalConnection = Connection(columnPoolerRegionInit['proxiamConnParams'])
        self.proximalConnection.createOneSegmentPerCell()
        self.internalDistalConnection = Connection(columnPoolerRegionInit['distalConnParams'])
        self.internalDistalConnection.createOneSegmentPerCell()
        self.lateralDistalConnection = {}
        for lc in lateralCorColOrders:
            self.lateralDistalConnection[lc] = Connection(columnPoolerRegionInit['distalConnParams'])
            self.lateralDistalConnection[lc].createOneSegmentPerCell()
        # Mutable States
        self.activeCells = np.array([])

    def reset(self):
        self.activeCells = np.array([])

    def compute(self, feedforwardInputs, lateralInputsMap, feedforwardGrowthCandidates):
        if not self.learn:
            self.computeInferenceMode(feedforwardInputs, lateralInputsMap)
        else:
            self.computeLearningMode(feedforwardInputs, lateralInputsMap, feedforwardGrowthCandidates)

    def computeInferenceMode(self, feedforwardInputs, lateralInputsMap):
        prevActiveCells = self.activeCells
        # Calculate the feedforward supported cells
        connectedOverlaps, _ = self.proximalConnection.computeActivity(feedforwardInputs)
        activeSegmentIds = self.proximalConnection.filterOverlapByThreshold(connectedOverlaps, self.proximalConnection.minThreshold)
        feedforwardSupportedCells = self.proximalConnection.mapSegmentsToCells(activeSegmentIds)
        # Calculate the number of active segments on each cell
        numActiveSegmentsByCell = np.zeros(self.cellCount, dtype="int")
        connectedOverlaps, _ = self.internalDistalConnection.computeActivity(prevActiveCells)
        activeSegmentIds = self.internalDistalConnection.getActiveSegments(connectedOverlaps)
        internaDistalSupportedCells = self.internalDistalConnection.mapSegmentsToCells(activeSegmentIds)
        numActiveSegmentsByCell[internaDistalSupportedCells] += 1
        for lcco in lateralInputsMap.keys():
            connectedOverlaps, _ = self.lateralDistalConnection[lcco].computeActivity(lateralInputsMap[lcco])
            activeSegmentIds = self.lateralDistalConnection[lcco].getActiveSegments(connectedOverlaps)
            laterDistalSupportedCells = self.lateralDistalConnection[lcco].mapSegmentsToCells(activeSegmentIds)
            numActiveSegmentsByCell[laterDistalSupportedCells] += 1
        # Calculate Chosen Cells
        chosenCells = []
        if len(feedforwardSupportedCells) == 0:
            pass
        else:
            numActiveSegsForFFSuppCells = numActiveSegmentsByCell[feedforwardSupportedCells]
            ttop = np.max(numActiveSegsForFFSuppCells)
            while ttop > 0 and len(chosenCells) < self.sdrSize:
                chosenCells = np.union1d(chosenCells, feedforwardSupportedCells[numActiveSegsForFFSuppCells >= ttop])
                ttop -= 1
        discrepancy = self.sdrSize - len(chosenCells)
        if discrepancy > 0:
            remFFcells = np.setdiff1d(feedforwardSupportedCells, chosenCells)
            n = (len(remFFcells) * discrepancy) // self.sdrSize
            n = max(n, discrepancy)
            n = min(n, len(remFFcells))
            if len(remFFcells) > n:
                selected = np.random.permutation(remFFcells)[0:n]
                chosenCells = np.append(chosenCells, selected)
            else:
                chosenCells = np.append(chosenCells, remFFcells)
        chosenCells.sort()
        self.activeCells = np.asarray(chosenCells, dtype="uint32")

    def computeLearningMode(self, feedforwardInputs, lateralInputsMap, feedforwardGrowthCandidates):
        prevActiveCells = self.activeCells
        if len(self.activeCells) < self.minSdrSize:
            self.activeCells = np.random.permutation(self.cellCount)[0:self.sdrSize]
            self.activeCells.sort()
        if len(feedforwardInputs) > 0:
            # Proximal
            learningSegmentIds = np.array([self.proximalConnection.allCells[cellId].segments[0].segmentId for cellId in self.activeCells])
            self.proximalConnection.learnOnExisting(learningSegmentIds, feedforwardInputs, feedforwardGrowthCandidates)
            # Internal
            if len(prevActiveCells) > 0:
                self.internalDistalConnection.learnOnExisting(learningSegmentIds, prevActiveCells, prevActiveCells)
            # Lateral
            for lcco in lateralInputsMap.keys():
                self.lateralDistalConnection[lcco].learnOnExisting(learningSegmentIds, lateralInputsMap[lcco], lateralInputsMap[lcco])
